#pragma once
#ifndef OBR_SHADER_DIRECTLIGHTSAMPLER_CUH
#define OBR_SHADER_DIRECTLIGHTSAMPLER_CUH

#include "../common.cuh"
#include "IntersectionTest.cuh"

namespace OBR
{
static __forceinline__ __device__ Spectrum shadeOneLight(SamplePRD *prd, LightInfo &desc)
{
#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("get light info: \n");
        printf("id: %d, instance: %d, type: %d\n", desc.light_id, desc.light_instance_id,
               desc.type);
    }
#endif
    const uint mat_base = prd->params->material_base_offset;

    Spectrum ld = Spectrum(0.0f);
    float4 urand = curand_uniform4(prd->rng_long);

    /* --------------------------- shade with MIS -------------------------- */
    float light_pdf = 1, surface_pdf = 1;
    Spectrum li, fr, tr;
    // light pass
    li = optixDirectCall<Spectrum, SamplePRD *, float *, int, float, float, LightCallMode>(
        desc.light_id + 1, prd, &light_pdf, desc.light_instance_id, urand.x, urand.y,
        LightCallMode::SAMPLE_LI);
    if (li.isValid())
    {
        SamplePRD occlusion_prd = *prd;
        occlusion_prd.ray_occlusion_test = true;
        intersectTr(&occlusion_prd, tr);
        if (!occlusion_prd.hit && light_pdf != 0)
        {
            fr = optixDirectCall<Spectrum, SamplePRD *, float *, float, float, MaterialCallMode>(
                prd->inter.mat_id + mat_base, prd, &surface_pdf, 0, 0,
                MaterialCallMode::CALC_FR_PDF);
            ld = fr * li * tr / light_pdf;
        }
    }

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("Ld for only light samping: \n");
        ld.print();
        printf("\n");
        printf("light pdf: %.6f, material pdf: %.6f\n", light_pdf, surface_pdf);
        printf("occlusion ray:\n");
        prd->inter.wi.print();
    }
#endif

    // surface pass
    if (desc.type != LightType::DELTA)
    {
        // add light weight to previous Ld
        float sqr_light_pdf = light_pdf * light_pdf;
        float sqr_surface_pdf = surface_pdf * surface_pdf;
        ld *= sqr_light_pdf / (sqr_light_pdf + sqr_surface_pdf);

        fr = optixDirectCall<Spectrum, SamplePRD *, float *, float, float, MaterialCallMode>(
            prd->inter.mat_id + mat_base, prd, &surface_pdf, urand.z, urand.w,
            MaterialCallMode::SAMPLE_FR);
        if (li.isValid())
        {
            SamplePRD occlusion_prd = *prd;
            occlusion_prd.ray_occlusion_test = true;
            intersectTr(&occlusion_prd, tr);
            if (occlusion_prd.hit && occlusion_prd.inter.light_id == desc.light_id &&
                occlusion_prd.inter.light_instance_id == desc.light_instance_id)
            {
                li = optixDirectCall<Spectrum, SamplePRD *, float *, int, float, float,
                                     LightCallMode>(desc.light_id + 1, prd, &light_pdf,
                                                    desc.light_instance_id, 0, 0,
                                                    LightCallMode::CALC_LI_PDF);
                sqr_light_pdf = light_pdf * light_pdf;
                sqr_surface_pdf = surface_pdf * surface_pdf;
                float weight = sqr_surface_pdf / (sqr_light_pdf + sqr_surface_pdf);
                ld += weight * fr * li * tr / surface_pdf;
            }
        }
    }

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("final Ld: \n");
        ld.print();
        printf("\n");
        printf("light pdf: %.6f, material pdf: %.6f\n", light_pdf, surface_pdf);
        printf("occlusion ray:\n");
        prd->inter.wi.print();
    }
#endif

    return ld;
}

static __forceinline__ __device__ Spectrum uniformSampleOneLight(SamplePRD *prd, float u0)
{
    int n_lights = prd->params->num_total_lights;
    if (n_lights > 0)
    {
        int light_desc_id = u0 * n_lights;
        light_desc_id = clamp(light_desc_id, 0, n_lights - 1);
        return shadeOneLight(prd, prd->params->light_desc[light_desc_id]) * n_lights;
    }
    else
    {
        return Spectrum(0.0f);
    }
}

static __forceinline__ __device__ Spectrum sampleAllLight(SamplePRD *prd)
{
    Spectrum L = Spectrum(0.0f);
    int n_lights = prd->params->num_total_lights;
    for (size_t light_desc_id = 0; light_desc_id < n_lights; light_desc_id++)
    {
        L += shadeOneLight(prd, prd->params->light_desc[light_desc_id]);
    }
    return L;
}
} // namespace OBR

#endif
